from crawler_base import *
from threading import BoundedSemaphore
from pytc import BDB, BDBOWRITER, BDBOCREAT, BDBTBZIP
from pytc import HDB, HDBOWRITER, HDBOCREAT, HDBTBZIP
import logging
from time import time
from math import floor
#import pymongo

class TmpTcDb(object):
    def __init__(self, filename):
        self.db = HDB()
        self.db.open(filename, HDBOCREAT | HDBOWRITER | HDBTBZIP)
        self.db.vanish()

    def put(self, key, value):
        self.db.put(key, value)

    def get(self, key):
        return self.get(key)

    def has_key(self, key):
        return self.db.has_key(key)

    def close(self):
        self.db.vanish()
        self.db.close()

class CrawlerDbTcIndexed(CrawlerDbInterface):
    def __init__(self, filename, maxdownload=None, cache=0):
        self.sem = BoundedSemaphore(2)
        self.filename=filename
        self.db = BDB()
        self.db.setmutex()
        leaf_cache=int(cache*(2/3.0))
        nonleaf_cache=int(cache*(1/3.0))
        self.db.setcache(leaf_cache,nonleaf_cache)
        self.db.open(self.filename, BDBOCREAT | BDBOWRITER | BDBTBZIP)
        #self.db.open(filename, BDBOCREAT | BDBOWRITER )
        self.maxdownload=maxdownload

	self.invalid_k='in'
        self.tovisit_k='tv'
        self.id_k='id'
        self.map_keytoid_k='ti'
        self.edges_k='e'
        self.visited_k='v'
        self.n=0
        self.edges_num=0
        self.visited_num=0
        self.init_stats=True
        self._setupDb()

    def __del__(self):
        self.db.close()

    def _setupDb(self):
        self._recover()
        self._optimize()
        self._sync()

    def _optimize(self):
        logging.info('Start database optimization')
        start=time()
        self.db.optimize(lmemb=0,nmemb=0,bnum=0,apow=-1,fpow=-1,opts=0xff)
        stop=time()
        logging.info('Database optimization took %s seconds', (stop-start))

    def _tovisitCacheCursor(self, db, key):
        counter = 0
        cursor = self.db.curnew()
        try:
            cursor.jump(key)
            while cursor.key().startswith(key):
                db.put(cursor.val(),'')
                cursor.next()
        except KeyError:
            pass

    def _recover(self):
        logging.info('Recovering previusly in visit keys... this may take some minutes')
        recovered=0
        start=time()
        tmpdb=  TmpTcDb(self.filename+'_rtmp')
        self._tovisitCacheCursor(tmpdb,self.tovisit_k)

        self.db.tranbegin()

        cursor=self.db.curnew()
        
        #keytoid_keys = self.db.rangefwm(self.map_keytoid_k, -1)
        try:
            cursor.jump(self.map_keytoid_k)
            while cursor.key().startswith(self.map_keytoid_k):
                if not self.db.has_key(self.visited_k+'_'+cursor.val()):
                    true_key = '_'.join(cursor.key().split('_')[1:])
                    if not tmpdb.has_key(true_key):
                        self._addToVisit(true_key)
                        recovered+=1
                cursor.next()
        except KeyError:
            pass

        self.db.trancommit()
        tmpdb.close()
        stop=time()
        logging.info('Recovered %d keys in %s seconds', recovered, (stop-start))

    def _sync(self):
        logging.info('Flushing data to disk')
        start=time()
        self.db.sync()
        stop=time()
        logging.info('Flush operation took %s seconds', (stop-start))        

    def optimize(self):
        self._optimize()

    def flush(self):
        self._sync()
        
    def printStats(self):
        logging.info('Processing statistics')
        start=time()
        res = self._calcStats()
        stop=time()
        logging.info('Statistics processed in %s', (stop-start))
        logging.info('Stats: tv=%s v=%s e=%s i=%s tk=%s',
                     res['tovisit'], res['visited'], res['edges'],
                     res['invalid'], res['totkeys'])

    def _visitedCounter(self, key):
        counter = 0
        cursor = self.db.curnew()
        try:
            cursor.jump(key)
            while cursor.key().startswith(key):
                counter+=1
                cursor.next()
        except KeyError:
            pass

        return counter

    def _edgesCounter(self, key):
        counter = 0
        cursor = self.db.curnew()
        try:
            cursor.jump(key)
            while cursor.key().startswith(key):
                idA,idB = cursor.key().split('_')[1:3]
                if idA < idB:
                    counter+=1
                cursor.next()
        except KeyError:
            pass

        return counter

    def _calcStats(self):
        res={}
        res['tovisit'] = self.db.vnum(self.tovisit_k)
        res['invalid'] = self.db.vnum(self.invalid_k)
        if self.init_stats:
            self.visited_num = self._visitedCounter(self.visited_k)
            self.edges_num = self._edgesCounter(self.edges_k)
            self.init_stats=False
        res['visited'] = self.visited_num
        res['edges'] = self.edges_num
        res['totkeys'] = 2*(res['tovisit']+res['visited'])+2*res['edges']+res['invalid']+1
        return res

    def _maxReached(self, n):
        if self.maxdownload == None:
            return False
        return n >= self.maxdownload

    def _getNextId(self):
        return str(self.db.addint(self.id_k, 1))

    def _getMapITKKey(self, key):
        return self.map_keytoid_k+'_'+key

    def _getIdFromKey(self, key):
        return str(self.db.get(self._getMapITKKey(key)))

    def _getEdgeKey(self, key, out_key):
        key_id = self._getIdFromKey(key)
        out_key_id = self._getIdFromKey(out_key)
        return '_'.join([self.edges_k,key_id,out_key_id]), \
               '_'.join([self.edges_k,out_key_id,key_id])

    def _getVisitedKey(self, key):
        key_id = self._getIdFromKey(key)
        return self.visited_k+'_'+key_id

    def _addToVisit(self, key):
            self.db.putdup(self.tovisit_k,key)

    def getToVisit(self):
        if self._maxReached(self.n):
            return None
        with self.sem:
            key=None
            if self.db.has_key(self.tovisit_k):
                key=self.db.get(self.tovisit_k)
                self.db.out(self.tovisit_k)
                return key

    def addToVisit(self, key):
        #if self._maxReached(self.n):
        #    return
        with self.sem:
            mapitk_key = self._getMapITKKey(key)
            if not self.db.has_key(mapitk_key):
                key_id = self._getNextId()
                self.db.putkeep(mapitk_key, key_id)
                self._addToVisit(key)

    ## def isVisited(self, key):
    ##     if not self.db.has_key(self.map_keytoid_k+key):
    ##         return False
    ##     key_id = self._getIdFromKey(key)
    ##     return self.db.has_key(self.edges_k+key_id)

    def addVisited(self, key, data):
        #if self._maxReached(self.n):
        #    return
        with self.sem:
            self.db.putkeep(self._getVisitedKey(key), data)
            self.n+=1
            self.visited_num+=1

    def addEdge(self, key, out_key, data):
        #if self._maxReached(self.n):
        #    return
        #edge_keyA = self._getEdgeKey(key, out_key)
        #edge_keyB = self._getEdgeKey(out_key, key)
        update_stats=False
        with self.sem:
            edge_keyA,edge_keyB = self._getEdgeKey(out_key, key)
            if not self.db.has_key(edge_keyA):
                self.db.put(edge_keyA,data)
                update_stats=True
            if not self.db.has_key(edge_keyB):            
                self.db.put(edge_keyB,data)
                update_stats=True
            if update_stats:
                self.edges_num+=1

    def invalidData(self, key):
        logging.warning('Invalid key (%s)', key)
        with self.sem:
            self.db.putdup(self.invalid_k, key)

    def toRetry(self, key):
        with self.sem:
            self._addToVisit(key)

class CrawlerDbMongo(CrawlerDbInterface):
    def __init__(self, filename, maxdownload=None, cache=0):
        self.connection = pymongo.Connection()
        self.db = self.connection[filename]
        self.connection.drop_database(self.db)
        #self.db.set_profiling_level(pymongo.ALL)
        self.maxdownload=maxdownload
        self.ttl_index_max=60*60*24*30
	self.invalid_k='in'
        self.tovisit_k='tv'
        self.id_k='id'
        self.map_keytoid_k='mkti'
        self.edges_k='e'
        self.visited_k='v'
        self.tovisit=self.db[self.tovisit_k]
        self.id = self.db[self.id_k]
        if self.id.count() == 0:
            self.id.insert({'_id':'keys', 'seq':1})
        self.map_keytoid=self.db[self.map_keytoid_k]
        self.visited=self.db[self.visited_k]
        self.edges=self.db[self.edges_k]
        self.n=0
        self._setupDb()

    def __del__(self):
        self.connection.disconnect()

    def _setupDb(self):
        self._recover()
        self._optimize()
        self._sync()

    def _optimize(self):
        logging.info('Start database optimization')
        start=time()
        self.tovisit.ensure_index([('_id', pymongo.ASCENDING)],ttl=self.ttl_index_max)
        self.map_keytoid.ensure_index([('_id', pymongo.ASCENDING)],ttl=self.ttl_index_max)
        self.map_keytoid.ensure_index([('key', pymongo.ASCENDING)],ttl=self.ttl_index_max)
        self.visited.ensure_index([('_id', pymongo.ASCENDING)],ttl=self.ttl_index_max)
        self.visited.ensure_index([('key', pymongo.ASCENDING)],ttl=self.ttl_index_max)
        self.edges.ensure_index([('_id', pymongo.ASCENDING)],ttl=self.ttl_index_max)
        self.edges.ensure_index([('keyA', pymongo.ASCENDING)],ttl=self.ttl_index_max)
        self.edges.ensure_index([('keyB', pymongo.ASCENDING)],ttl=self.ttl_index_max)
        self.edges.ensure_index([('keyA', pymongo.ASCENDING),('keyB', pymongo.ASCENDING)],
                                ttl=self.ttl_index_max,unique=True,drop_dups=True)
        self.edges.ensure_index([('_id', pymongo.ASCENDING),('keyB', pymongo.ASCENDING)],
                                ttl=self.ttl_index_max,unique=True,drop_dups=True)
        stop=time()
        logging.info('Database optimization took %s seconds', (stop-start))

    def _recover(self):
        logging.info('Recovering previusly in visit keys... this may take some minutes')
        recovered=0
        start=time()
        stop=time()
        logging.info('Recovered %d keys in %s seconds', recovered, (stop-start))

    def _sync(self):
        logging.info('Flushing data to disk')
        start=time()
        self.connection.fsync()
        stop=time()
        logging.info('Flush operation took %s seconds', (stop-start))        
    def optimize(self):
        self._optimize()

    def flush(self):
        self._sync()
        
    def printStats(self):
        logging.info('Processing statistics')
        start=time()
        res = self._calcStats()
        stop=time()
        logging.info('Statistics processed in %s', (stop-start))
        logging.info('Stats: tovisit=%s   visited=%s   edges=%s',res['tovisit'], res['visited'], res['edges'])

    def _calcStats(self):
        res={}
        res['tovisit'] = self.tovisit.count()
        res['edges'] = self.edges.count()
        res['visited'] = self.visited.count()

        return res

    def _maxReached(self, n):
        if self.maxdownload == None:
            return False
        return n >= self.maxdownload

    def _getNextId(self):
        return self.id.find_and_modify(query={'_id':'keys'},
                                       update={'$inc': {'seq':1}},
                                       new=True)

    def _getIdFromKey(self, key):
        return self.map_keytoid.find_one({'key':key})
    
    def getToVisit(self):
        if self._maxReached(self.n):
            return None        
        key=self.tovisit.find_one()
        if key <> None:
            self.tovisit.remove(key)
            key=key['key']
        return key

    def addToVisit(self, key):
        if self._maxReached(self.n):
            return
        if self.map_keytoid.find_one({'key':key}) == None:
            next_id = self._getNextId()
            self.map_keytoid.insert({'_id':next_id,'key':key})
            self.tovisit.insert({'_id':next_id,'key':key})

    def addVisited(self, key, data):
        if self._maxReached(self.n):
            return
        self.visited.insert({'_id':key,'data':data})
        self.n+=1

    def addEdge(self, keyA, keyB, data):
        if self._maxReached(self.n):
            return
        keyA_id=self._getIdFromKey(keyA)
        keyB_id=self._getIdFromKey(keyB)
        a_id,b_id = keyB_id,keyA_id
        if keyA_id < keyB_id:
            a_id,b_id = keyA_id,keyB_id
        self.edges.insert({'keyA':a_id,'keyB':b_id,'data':data})

    def invalidData(self, key):
        logging.warning('Invalid key (%s)', key)
        self.invalid.insert({'key':key})

    def toRetry(self, key):
        next_id = self._getNextId()
        self.tovisit.insert({'_id':next_id,'key':key})
